import torch
import sys
import os

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from model.attnSleepModel import AttnSleep
from util.webUtils import jsonToTensor

# 模型预测
def modelPredictServer(segmentData):
    # json转tensor
    segment_tensor=jsonToTensor(segmentData)

    
    model = AttnSleep()
    model.load_state_dict(torch.load('saved/trained_model.pth'))
    model.eval()
    with torch.no_grad():
        outputs = model(segment_tensor)
        _, predicted_class = torch.max(outputs, 1)
        predicted_stage = predicted_class.item()
        print("Predicted sleep stage for the segment:", predicted_stage)

        # 映射到标准标签，统一全局分期定义
        print("Mapped to standard label:", predicted_stage)
        return predicted_stage
    
    # 保留原始返回作为备选方案
    # return predicted_stage

 # def _valid_epoch(self, epoch):
 #        """
 #        Validate after training an epoch
 #
 #        :param epoch: Integer, current training epoch.
 #        :return: A log that contains information about validation
 #        """
 #        self.model.eval()
 #        self.valid_metrics.reset()
 #        with torch.no_grad():
 #            outs = np.array([])
 #            trgs = np.array([])
 #            for batch_idx, (data, target) in enumerate(self.valid_data_loader):
 #                data, target = data.to(self.device), target.to(self.device)
 #                output = self.model(data)
 #                loss = self.criterion(output, target, self.class_weights, self.device)
 #
 #                self.valid_metrics.update('loss', loss.item())
 #                for met in self.metric_ftns:
 #                    self.valid_metrics.update(met.__name__, met(output, target))
 #
 #                preds_ = output.data.max(1, keepdim=True)[1].cpu()
 #
 #                outs = np.append(outs, preds_.cpu().numpy())
 #                trgs = np.append(trgs, target.data.cpu().numpy())
 #
 #
 #        return self.valid_metrics.result(), outs, trgs
 #
